import requests
import json
import time
import os
import logging
from dotenv import load_dotenv

# 加载 .env 文件
load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 配置
COMFYUI_SERVER = os.getenv('COMFYUI_SERVER', 'http://127.0.0.1:7860')
WEB_SERVER = os.getenv('WEB_SERVER', 'http://localhost:15000')
POLL_INTERVAL = int(os.getenv('POLL_INTERVAL', '5'))  # 轮询间隔（秒）
OUTPUT_FOLDER = os.getenv('OUTPUT_FOLDER', './outputs')  # 输出目录

def update_workflow_with_image(workflow_data, image_name):
    """更新工作流配置，替换输入图片"""
    # 查找 LoadImage 节点
    load_image_nodes = []
    for node_id, node in workflow_data.items():
        if node.get('class_type') == 'LoadImage':
            load_image_nodes.append(node_id)
    
    # 如果找到 LoadImage 节点，更新图片路径
    if load_image_nodes:
        for node_id in load_image_nodes:
            workflow_data[node_id]['inputs']['image'] = image_name
    else:
        # 如果没有 LoadImage 节点，创建一个并更新相关连接
        new_node_id = str(max(int(k) for k in workflow_data.keys()) + 1)
        workflow_data[new_node_id] = {
            "class_type": "LoadImage",
            "inputs": {
                "image": image_name
            }
        }
        # TODO: 根据需要添加与其他节点的连接

    return workflow_data

def wait_for_prompt(prompt_id, timeout=300):
    """等待工作流执行完成"""
    start_time = time.time()
    while True:
        if time.time() - start_time > timeout:
            raise Exception("Workflow execution timeout")
            
        try:
            response = requests.get(f"{COMFYUI_SERVER}/history/{prompt_id}")
            if response.status_code == 200:
                data = response.json()
                if prompt_id in data:
                    prompt_data = data[prompt_id]
                    logger.info(f"Prompt status: {json.dumps(prompt_data, indent=2)}")
                    
                    # 检查是否有错误
                    if "error" in prompt_data:
                        raise Exception(f"Workflow execution failed: {prompt_data['error']}")
                    
                    # 检查是否完成
                    if "outputs" in prompt_data and prompt_data.get("status", {}).get("completed", False):
                        return prompt_data
            
            # 等待一段时间再检查
            time.sleep(1)
            
        except Exception as e:
            logger.error(f"Error checking prompt status: {str(e)}")
            time.sleep(1)

def process_task(task_data):
    """处理单个任务"""
    try:
        # 下载图片
        image_path = task_data['image_path']
        logger.info(f"Downloading image from: {WEB_SERVER}/uploads/{os.path.basename(image_path)}")
        response = requests.get(f"{WEB_SERVER}/uploads/{os.path.basename(image_path)}")
        if response.status_code != 200:
            raise Exception("Failed to download image")
        
        local_image_path = os.path.basename(image_path)
        with open(local_image_path, 'wb') as f:
            f.write(response.content)
        logger.info(f"Image saved locally as: {local_image_path}")
        
        # 上传图片到ComfyUI
        logger.info(f"Uploading image to ComfyUI: {COMFYUI_SERVER}/upload/image")
        with open(local_image_path, 'rb') as f:
            response = requests.post(f"{COMFYUI_SERVER}/upload/image", files={'image': f})
            if response.status_code != 200:
                raise Exception("Failed to upload image to ComfyUI")
            image_name = response.json()['name']
            logger.info(f"Image uploaded to ComfyUI, name: {image_name}")
        
        # 更新工作流中的图片路径
        workflow_data = task_data['workflow']
        logger.info("Original workflow data: " + json.dumps(workflow_data, indent=2))
        
        # 确保工作流数据格式正确
        if not isinstance(workflow_data, dict):
            raise Exception("Invalid workflow data format")
            
        # 构建提交到 ComfyUI 的数据格式
        prompt_data = {
            "prompt": workflow_data,
            "client_id": f"worker_{int(time.time())}",
            "extra_data": {
                "image_name": image_name
            }
        }
        
        # 更新工作流中的图片路径
        workflow_data = update_workflow_with_image(workflow_data, image_name)
        logger.info("Updated workflow data: " + json.dumps(workflow_data, indent=2))
        
        # 提交任务到ComfyUI
        logger.info(f"Submitting workflow to ComfyUI: {COMFYUI_SERVER}/prompt")
        response = requests.post(
            f"{COMFYUI_SERVER}/prompt",
            json=prompt_data,
            headers={'Content-Type': 'application/json'}
        )
        logger.info(f"ComfyUI response status: {response.status_code}")
        logger.info(f"ComfyUI response content: {response.text}")
        
        if response.status_code != 200:
            raise Exception(f"Failed to submit workflow: {response.text}")
        
        prompt_id = response.json()['prompt_id']
        logger.info(f"Waiting for workflow execution (prompt_id: {prompt_id})")
        
        # 等待工作流执行完成
        result = wait_for_prompt(prompt_id)
        logger.info(f"Workflow execution completed: {json.dumps(result, indent=2)}")
        
        # 处理输出结果
        if "outputs" in result:
            outputs = result['outputs']
            # 处理输出图片
            processed_outputs = {}
            for node_id, output in outputs.items():
                if 'images' in output:
                    processed_outputs[node_id] = {
                        'images': []
                    }
                    for img_data in output['images']:
                        if 'filename' in img_data:
                            # 从 ComfyUI 获取图片
                            img_url = f"{COMFYUI_SERVER}/view?filename={img_data['filename']}"
                            response = requests.get(img_url)
                            if response.status_code == 200:
                                # 生成唯一的输出文件名
                                output_filename = f"output_{int(time.time())}_{img_data['filename']}"
                                
                                # 将图片发送给 ServerB
                                files = {
                                    'image': (output_filename, response.content, 'image/png')
                                }
                                upload_response = requests.post(
                                    f"{WEB_SERVER}/save_output",
                                    files=files
                                )
                                
                                if upload_response.status_code != 200:
                                    raise Exception("Failed to upload output image to ServerB")
                                
                                processed_outputs[node_id]['images'].append({
                                    'filename': output_filename
                                })
            
            # 将结果回传给 ServerB
            response = requests.post(
                f"{WEB_SERVER}/task/complete",
                json={
                    'task_id': task_data['task_id'],
                    'result': processed_outputs
                }
            )
            if response.status_code != 200:
                raise Exception("Failed to send result back")
        else:
            raise Exception("No outputs in workflow result")
            
        logger.info(f"Task {task_data['task_id']} completed successfully")
        
        # 清理临时文件
        try:
            os.remove(local_image_path)
        except Exception as e:
            logger.warning(f"Failed to clean up temporary file: {str(e)}")
            
    except Exception as e:
        logger.error(f"Error processing task {task_data['task_id']}: {str(e)}")
        # 通知任务失败
        try:
            requests.post(
                f"{WEB_SERVER}/task/complete",
                json={
                    'task_id': task_data['task_id'],
                    'error': str(e)
                }
            )
        except Exception as send_error:
            logger.error(f"Failed to send error status: {str(send_error)}")
            
        # 清理临时文件
        try:
            if 'local_image_path' in locals():
                os.remove(local_image_path)
        except Exception as cleanup_error:
            logger.warning(f"Failed to clean up temporary file: {str(cleanup_error)}")

def main():
    """主循环：不断从 ServerB 获取和处理任务"""
    logger.info("Worker started")
    
    while True:
        try:
            # 尝试获取任务
            response = requests.post(f"{WEB_SERVER}/task/pull")
            
            if response.status_code == 200:
                # 获取到任务，开始处理
                task_data = response.json()
                logger.info(f"Received task: {task_data['task_id']}")
                process_task(task_data)
            elif response.status_code == 404:
                # 没有任务，等待后重试
                logger.debug("No task available")
                time.sleep(POLL_INTERVAL)
            else:
                logger.error(f"Error pulling task: {response.status_code}")
                time.sleep(POLL_INTERVAL)
                
        except Exception as e:
            logger.error(f"Error in main loop: {str(e)}")
            time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
    main()
